import asyncio
import json
import base64
from pathlib import Path

async def process_tcp_data():
    # 初始化状态变量
    buffer = bytearray()
    banner_processed = False
    frame_body_length = 0
    frame_count = 0

    current_banner = {
        'version': 0,
        'length': 0,
        'pid': 0,
        'realWidth': 0,
        'realHeight': 0,
        'virtualWidth': 0,
        'virtualHeight': 0,
        'orientation': 0,
        'quirks': 0
    }

    # 创建输出目录
    output_dir = Path('output')
    output_dir.mkdir(parents=True, exist_ok=True)

    # 建立 TCP 连接
    reader, writer = await asyncio.open_connection('localhost', 1717)

    try:
        while True:
            # 读取 TCP 数据
            data = await reader.read(4096)
            if not data:
                break
            
            # 将数据添加到缓冲区
            buffer.extend(data)

            # 持续处理缓冲区中的数据
            while True:
                # 解析 Banner
                if not banner_processed:
                    if len(buffer) < 2:
                        break  # 等待更多数据

                    version = buffer[0]
                    banner_length = buffer[1]
                    current_banner['version'] = version
                    current_banner['length'] = banner_length

                    if len(buffer) < banner_length:
                        break  # 数据不足

                    banner_bytes = bytes(buffer[:banner_length])
                    buffer = buffer[banner_length:]

                    # 解析 Banner 字段（小端序）
                    current_banner['pid'] = int.from_bytes(banner_bytes[2:6], 'little')
                    current_banner['realWidth'] = int.from_bytes(banner_bytes[6:10], 'little')
                    current_banner['realHeight'] = int.from_bytes(banner_bytes[10:14], 'little')
                    current_banner['virtualWidth'] = int.from_bytes(banner_bytes[14:18], 'little')
                    current_banner['virtualHeight'] = int.from_bytes(banner_bytes[18:22], 'little')
                    current_banner['orientation'] = banner_bytes[22] * 90
                    current_banner['quirks'] = banner_bytes[23]

                    # 保存 Banner 信息
                    banner_info = {
                        "raw_data": base64.b64encode(banner_bytes).decode(),
                        "parsed": current_banner
                    }
                    with open('output/banner.json', 'w') as f:
                        json.dump(banner_info, f, indent=2)
                    
                    banner_processed = True
                    print("Banner processed")

                # 解析帧数据
                else:
                    if frame_body_length == 0:
                        if len(buffer) < 4:
                            break  # 需要更多数据
                        frame_body_length = int.from_bytes(buffer[:4], 'little')
                        buffer = buffer[4:]

                    if len(buffer) < frame_body_length:
                        break

                    frame_data = bytes(buffer[:frame_body_length])
                    buffer = buffer[frame_body_length:]

                    # 验证 JPEG 头
                    if frame_data[0] != 0xFF or frame_data[1] != 0xD8:
                        print("Invalid JPEG header")
                        frame_body_length = 0
                        continue

                    # 保存帧图片
                    frame_count += 1
                    Path(f'output/frame_{frame_count}.jpg').write_bytes(frame_data)
                    print(f"Saved frame {frame_count}")
                    frame_body_length = 0  # 重置状态

    except Exception as e:
        print(f"Error: {e}")
    finally:
        writer.close()
        print("Connection closed")

if __name__ == "__main__":
    asyncio.run(process_tcp_data())